from core.algorithm.deeplearning.mmdeploy.detectors.detector import MMDetector
from glob import glob
import cv2, os, json

model_dir = "/home/smartgis/workspace/project/mmlab/mmyolo/_mmyolo_custom/projects/eve/yolov8_m/deploy"
file_dir = "data/input"
save_dir = "data/output"

cls_map = {
    0:"bar",
    1:"list",
    2:"list_box",
    3:"target",
    4:"object",
    5:"button",
    6:"weapon",
    7:"mining",
    8:"buff",
    9:"right_mouse",
}

def building_labelme(height,width, base_name,  bboxes):
    labelme_json = {
                    "version": "5.0.2", 
                    "flags": {},
                    "shapes": [],
                    "imagePath": "../data/"+base_name,
                    "imageHeight": height,
                    "imageWidth": width,
                    "imageData":None,
            }
        # 找出当前文件的anno, 并且按照labelme的格式保存
    for bbox in bboxes:
        cls_idx, score, x1, y1, x2 ,y2 = bbox
        if cls_idx not in cls_map.keys():
            continue
        cls_name = cls_map[cls_idx]
        shape = {"label":cls_name,
                    "points":[[x1,y1],[x2,y2]],
                    "group_id": None,
                    "shape_type": "rectangle",
                    "flags":{}
                    }
        labelme_json["shapes"].append(shape)
    return labelme_json

detector = MMDetector(model_dir=model_dir,
                keys={
                    "in":"input_data",
                    "out":"detect",
                },
                device="cpu"
                )

file_list = list(glob(file_dir+"/*"))

for file_path in file_list:
    img_data = cv2.imread(file_path)[None,...]
    file_name, extension = os.path.basename(file_path).rsplit(".",1)
    _, height, width,_ = img_data.shape
    data = {
            "input_data":img_data
        }
    data = detector(data)
    labelme_json = building_labelme(width, height, os.path.basename(file_path), data["detect"]["boxes"][0])
    with open(save_dir+"/{}.json".format(file_name),"w") as f:
        json.dump(labelme_json,f)


